using System;
using System.Collections.Generic;
using System.Linq;
using UnityEngine;

namespace GenerativeAgents.Core
{
    /// <summary>
    /// 联想记忆系统
    /// 存储和管理智能体的所有记忆节点，支持基于语义相似度的检索
    /// </summary>
    [System.Serializable]
    public class AssociativeMemory
    {
        [Header("联想记忆配置")]
        public int maxCapacity = 1000;
        public float forgettingThreshold = 0.1f; // 遗忘阈值
        public int minAccessCountForRetention = 3; // 保留所需的最小访问次数
        
        // 记忆存储
        private List<MemoryNode> memoryNodes;
        private Dictionary<string, MemoryNode> memoryIndex; // 快速查找索引
        private Dictionary<MemoryType, List<MemoryNode>> typeIndex; // 类型索引
        
        // 记忆统计
        private int totalMemoriesAdded = 0;
        private int totalMemoriesRemoved = 0;
        private Dictionary<MemoryType, int> typeStatistics;
        
        // 遗忘机制
        private DateTime lastForgetCycle = DateTime.Now;
        private TimeSpan forgetCycleInterval = TimeSpan.FromMinutes(10); // 10分钟执行一次遗忘检查
        
        public AssociativeMemory(int capacity = 1000)
        {
            maxCapacity = capacity;
            Initialize();
        }
        
        /// <summary>
        /// 初始化联想记忆系统
        /// </summary>
        private void Initialize()
        {
            memoryNodes = new List<MemoryNode>();
            memoryIndex = new Dictionary<string, MemoryNode>();
            typeIndex = new Dictionary<MemoryType, List<MemoryNode>>();
            typeStatistics = new Dictionary<MemoryType, int>();
            
            // 初始化类型索引
            foreach (MemoryType type in Enum.GetValues(typeof(MemoryType)))
            {
                typeIndex[type] = new List<MemoryNode>();
                typeStatistics[type] = 0;
            }
        }
        
        /// <summary>
        /// 添加记忆节点
        /// </summary>
        public void AddMemory(MemoryNode memoryNode)
        {
            if (memoryNode == null || string.IsNullOrEmpty(memoryNode.nodeId))
            {
                Debug.LogWarning("[AssociativeMemory] Cannot add null or invalid memory node");
                return;
            }
            
            try
            {
                // 检查是否已存在
                if (memoryIndex.ContainsKey(memoryNode.nodeId))
                {
                    Debug.LogWarning($"[AssociativeMemory] Memory node {memoryNode.nodeId} already exists");
                    return;
                }
                
                // 检查容量并执行遗忘机制
                if (memoryNodes.Count >= maxCapacity)
                {
                    ForgetOldMemories();
                }
                
                // 添加记忆
                memoryNodes.Add(memoryNode);
                memoryIndex[memoryNode.nodeId] = memoryNode;
                typeIndex[memoryNode.type].Add(memoryNode);
                
                // 更新统计
                totalMemoriesAdded++;
                typeStatistics[memoryNode.type]++;
                
                // 建立关联
                EstablishAssociations(memoryNode);
                
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error adding memory {memoryNode.nodeId}: {ex.Message}");
            }
        }
        
        /// <summary>
        /// 移除记忆节点
        /// </summary>
        public bool RemoveMemory(string nodeId)
        {
            if (!memoryIndex.TryGetValue(nodeId, out var memoryNode))
            {
                return false;
            }
            
            try
            {
                // 从各种索引中移除
                memoryNodes.Remove(memoryNode);
                memoryIndex.Remove(nodeId);
                typeIndex[memoryNode.type].Remove(memoryNode);
                
                // 移除关联
                RemoveAssociations(memoryNode);
                
                // 更新统计
                totalMemoriesRemoved++;
                typeStatistics[memoryNode.type]--;
                
                return true;
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error removing memory {nodeId}: {ex.Message}");
                return false;
            }
        }
        
        /// <summary>
        /// 获取记忆节点
        /// </summary>
        public MemoryNode GetMemory(string nodeId)
        {
            memoryIndex.TryGetValue(nodeId, out var memory);
            return memory;
        }
        
        /// <summary>
        /// 搜索相似记忆
        /// </summary>
        public List<MemoryNode> SearchSimilar(float[] queryEmbedding, int topK = 10, MemoryType? typeFilter = null)
        {
            if (queryEmbedding == null)
            {
                return new List<MemoryNode>();
            }
            
            try
            {
                var searchNodes = typeFilter.HasValue ? typeIndex[typeFilter.Value] : memoryNodes;
                var similarities = new List<(MemoryNode node, float similarity)>();
                
                foreach (var node in searchNodes)
                {
                    if (node.CheckExpired()) continue;
                    
                    if (node.embedding != null && node.embedding.Length == queryEmbedding.Length)
                    {
                        var similarity = CalculateCosineSimilarity(queryEmbedding, node.embedding);
                        similarities.Add((node, similarity));
                    }
                }
                
                return similarities
                    .OrderByDescending(x => x.similarity)
                    .Take(topK)
                    .Select(x => x.node)
                    .ToList();
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error searching similar memories: {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 基于关键词搜索
        /// </summary>
        public List<MemoryNode> SearchByKeywords(List<string> keywords, int topK = 10, MemoryType? typeFilter = null)
        {
            if (keywords == null || keywords.Count == 0)
            {
                return new List<MemoryNode>();
            }
            
            try
            {
                var searchNodes = typeFilter.HasValue ? typeIndex[typeFilter.Value] : memoryNodes;
                var scores = new List<(MemoryNode node, float score)>();
                
                foreach (var node in searchNodes)
                {
                    if (node.CheckExpired()) continue;
                    
                    var score = CalculateKeywordScore(node, keywords);
                    if (score > 0)
                    {
                        scores.Add((node, score));
                    }
                }
                
                return scores
                    .OrderByDescending(x => x.score)
                    .Take(topK)
                    .Select(x => x.node)
                    .ToList();
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error searching by keywords: {ex.Message}");
                return new List<MemoryNode>();
            }
        }
        
        /// <summary>
        /// 获取指定类型的记忆
        /// </summary>
        public List<MemoryNode> GetMemoriesByType(MemoryType type, int maxCount = -1)
        {
            var memories = typeIndex[type].Where(m => !m.CheckExpired()).ToList();
            
            if (maxCount > 0 && memories.Count > maxCount)
            {
                // 按重要性和新近度排序
                memories = memories
                    .OrderByDescending(m => m.GetCompositeImportance())
                    .ThenByDescending(m => m.lastAccessTime)
                    .Take(maxCount)
                    .ToList();
            }
            
            return memories;
        }
        
        /// <summary>
        /// 获取最近的记忆
        /// </summary>
        public List<MemoryNode> GetRecentMemories(TimeSpan timeSpan, int maxCount = 20)
        {
            var cutoffTime = DateTime.Now - timeSpan;
            
            return memoryNodes
                .Where(m => m.createdTime >= cutoffTime && !m.CheckExpired())
                .OrderByDescending(m => m.createdTime)
                .Take(maxCount)
                .ToList();
        }
        
        /// <summary>
        /// 获取高重要性记忆
        /// </summary>
        public List<MemoryNode> GetHighImportanceMemories(int minImportance = 7, int maxCount = 15)
        {
            return memoryNodes
                .Where(m => m.poignancy >= minImportance && !m.CheckExpired())
                .OrderByDescending(m => m.poignancy)
                .ThenByDescending(m => m.accessCount)
                .Take(maxCount)
                .ToList();
        }
        
        /// <summary>
        /// 获取相关记忆
        /// </summary>
        public List<MemoryNode> GetRelatedMemories(string nodeId, int maxCount = 10)
        {
            if (!memoryIndex.TryGetValue(nodeId, out var sourceNode))
            {
                return new List<MemoryNode>();
            }
            
            var relatedMemories = new List<MemoryNode>();
            
            // 添加直接关联的记忆
            foreach (var relatedId in sourceNode.relatedNodeIds)
            {
                if (memoryIndex.TryGetValue(relatedId, out var relatedNode) && !relatedNode.CheckExpired())
                {
                    relatedMemories.Add(relatedNode);
                }
            }
            
            // 如果关联记忆不够，基于语义相似度添加更多
            if (relatedMemories.Count < maxCount && sourceNode.embedding != null)
            {
                var similarMemories = SearchSimilar(sourceNode.embedding, maxCount * 2)
                    .Where(m => m.nodeId != nodeId && !relatedMemories.Contains(m))
                    .Take(maxCount - relatedMemories.Count);
                
                relatedMemories.AddRange(similarMemories);
            }
            
            return relatedMemories.Take(maxCount).ToList();
        }
        
        /// <summary>
        /// 遗忘旧记忆
        /// </summary>
        private void ForgetOldMemories()
        {
            try
            {
                // 找到最不重要和最少访问的记忆
                var forgettableMemories = memoryNodes
                    .Where(m => m.accessCount < minAccessCountForRetention || m.GetRecencyScore() < forgettingThreshold)
                    .OrderBy(m => m.GetCompositeImportance())
                    .ThenBy(m => m.accessCount)
                    .Take(maxCapacity / 10) // 每次遗忘10%的容量
                    .ToList();
                
                foreach (var memory in forgettableMemories)
                {
                    RemoveMemory(memory.nodeId);
                }
                
                Debug.Log($"[AssociativeMemory] Forgot {forgettableMemories.Count} memories to make space");
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error during forgetting process: {ex.Message}");
            }
        }
        
        /// <summary>
        /// 建立记忆关联
        /// </summary>
        private void EstablishAssociations(MemoryNode newMemory)
        {
            try
            {
                // 基于时间邻近性建立关联
                var recentMemories = GetRecentMemories(TimeSpan.FromHours(1), 10);
                foreach (var recentMemory in recentMemories)
                {
                    if (recentMemory.nodeId != newMemory.nodeId && ShouldEstablishAssociation(newMemory, recentMemory))
                    {
                        newMemory.AddRelatedMemory(recentMemory.nodeId);
                        recentMemory.AddRelatedMemory(newMemory.nodeId);
                    }
                }
                
                // 基于参与者建立关联
                if (newMemory.participants != null)
                {
                    foreach (var participant in newMemory.participants)
                    {
                        var participantMemories = memoryNodes
                            .Where(m => m.participants != null && m.participants.Contains(participant) && m.nodeId != newMemory.nodeId)
                            .OrderByDescending(m => m.createdTime)
                            .Take(5);
                        
                        foreach (var memory in participantMemories)
                        {
                            newMemory.AddRelatedMemory(memory.nodeId);
                            memory.AddRelatedMemory(newMemory.nodeId);
                        }
                    }
                }
                
                // 基于关键词建立关联
                if (newMemory.keywords != null)
                {
                    var keywordMemories = SearchByKeywords(newMemory.keywords, 5)
                        .Where(m => m.nodeId != newMemory.nodeId);
                    
                    foreach (var memory in keywordMemories)
                    {
                        newMemory.AddRelatedMemory(memory.nodeId);
                        memory.AddRelatedMemory(newMemory.nodeId);
                    }
                }
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error establishing associations for {newMemory.nodeId}: {ex.Message}");
            }
        }
        
        /// <summary>
        /// 移除记忆关联
        /// </summary>
        private void RemoveAssociations(MemoryNode memoryToRemove)
        {
            try
            {
                foreach (var relatedId in memoryToRemove.relatedNodeIds)
                {
                    if (memoryIndex.TryGetValue(relatedId, out var relatedMemory))
                    {
                        relatedMemory.relatedNodeIds.Remove(memoryToRemove.nodeId);
                    }
                }
            }
            catch (Exception ex)
            {
                Debug.LogError($"[AssociativeMemory] Error removing associations for {memoryToRemove.nodeId}: {ex.Message}");
            }
        }
        
        /// <summary>
        /// 判断是否应该建立关联
        /// </summary>
        private bool ShouldEstablishAssociation(MemoryNode memory1, MemoryNode memory2)
        {
            // 同一位置
            if (Vector3.Distance(memory1.location, memory2.location) < 5f)
                return true;
            
            // 共同参与者
            if (memory1.participants != null && memory2.participants != null)
            {
                var commonParticipants = memory1.participants.Intersect(memory2.participants).Count();
                if (commonParticipants > 0)
                    return true;
            }
            
            // 相似关键词
            if (memory1.keywords != null && memory2.keywords != null)
            {
                var commonKeywords = memory1.keywords.Intersect(memory2.keywords).Count();
                if (commonKeywords >= 2)
                    return true;
            }
            
            // 语义相似度
            if (memory1.embedding != null && memory2.embedding != null)
            {
                var similarity = CalculateCosineSimilarity(memory1.embedding, memory2.embedding);
                if (similarity > 0.7f)
                    return true;
            }
            
            return false;
        }
        
        /// <summary>
        /// 计算余弦相似度
        /// </summary>
        private float CalculateCosineSimilarity(float[] vector1, float[] vector2)
        {
            if (vector1.Length != vector2.Length)
                return 0f;
            
            float dotProduct = 0f;
            float norm1 = 0f;
            float norm2 = 0f;
            
            for (int i = 0; i < vector1.Length; i++)
            {
                dotProduct += vector1[i] * vector2[i];
                norm1 += vector1[i] * vector1[i];
                norm2 += vector2[i] * vector2[i];
            }
            
            if (norm1 == 0f || norm2 == 0f)
                return 0f;
            
            return dotProduct / (Mathf.Sqrt(norm1) * Mathf.Sqrt(norm2));
        }
        
        /// <summary>
        /// 计算关键词匹配分数
        /// </summary>
        private float CalculateKeywordScore(MemoryNode node, List<string> queryKeywords)
        {
            if (node.keywords == null || node.keywords.Count == 0)
                return 0f;
            
            var matchCount = 0f;
            foreach (var keyword in queryKeywords)
            {
                foreach (var nodeKeyword in node.keywords)
                {
                    if (string.Equals(keyword, nodeKeyword, StringComparison.OrdinalIgnoreCase))
                    {
                        matchCount += 1f;
                    }
                    else if (nodeKeyword.Contains(keyword, StringComparison.OrdinalIgnoreCase) ||
                             keyword.Contains(nodeKeyword, StringComparison.OrdinalIgnoreCase))
                    {
                        matchCount += 0.5f;
                    }
                }
            }
            
            return matchCount / Math.Max(queryKeywords.Count, node.keywords.Count);
        }
        
        /// <summary>
        /// 执行定期维护
        /// </summary>
        public void PerformMaintenance()
        {
            if (DateTime.Now - lastForgetCycle > forgetCycleInterval)
            {
                // 清理过期记忆
                var expiredMemories = GetExpiredMemories();
                foreach (var memory in expiredMemories)
                {
                    RemoveMemory(memory.nodeId);
                }
                
                // 如果还是超容量，执行遗忘
                if (memoryNodes.Count > maxCapacity)
                {
                    ForgetOldMemories();
                }
                
                lastForgetCycle = DateTime.Now;
            }
        }
        
        /// <summary>
        /// 获取过期记忆
        /// </summary>
        public List<MemoryNode> GetExpiredMemories()
        {
            return memoryNodes.Where(m => m.CheckExpired()).ToList();
        }
        
        /// <summary>
        /// 获取所有记忆
        /// </summary>
        public List<MemoryNode> GetAllMemories()
        {
            return new List<MemoryNode>(memoryNodes);
        }
        
        /// <summary>
        /// 获取记忆数量
        /// </summary>
        public int GetMemoryCount()
        {
            return memoryNodes.Count;
        }
        
        /// <summary>
        /// 设置最大容量
        /// </summary>
        public void SetMaxCapacity(int capacity)
        {
            maxCapacity = capacity;
            if (memoryNodes.Count > maxCapacity)
            {
                ForgetOldMemories();
            }
        }
        
        /// <summary>
        /// 获取统计信息
        /// </summary>
        public string GetStatistics()
        {
            var stats = $"Total Memories: {memoryNodes.Count}/{maxCapacity}\n";
            stats += $"Added: {totalMemoriesAdded}, Removed: {totalMemoriesRemoved}\n";
            stats += "Type Distribution:\n";
            
            foreach (var kvp in typeStatistics)
            {
                if (kvp.Value > 0)
                {
                    stats += $"  {kvp.Key}: {kvp.Value}\n";
                }
            }
            
            return stats;
        }
        
        /// <summary>
        /// 清空所有记忆
        /// </summary>
        public void ClearAllMemories()
        {
            memoryNodes.Clear();
            memoryIndex.Clear();
            
            foreach (var type in typeIndex.Keys.ToList())
            {
                typeIndex[type].Clear();
                typeStatistics[type] = 0;
            }
            
            totalMemoriesAdded = 0;
            totalMemoriesRemoved = 0;
        }
    }
}